# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-09-25

import sys
from pathlib import Path
from tqdm import tqdm
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim

from Algorithms.Transformer_DL.transformer_model import TransformerModel
from Common.time_utils import generate_timestamps
from config import TaskConfig


current_file = Path(__file__).resolve()
parent_dir = current_file.parent
# 将父目录添加到 sys.path
sys.path.insert(0, str(parent_dir))
device = torch.device(TaskConfig.DEVICE if torch.cuda.is_available() else "cpu")


class Model(TransformerModel):
    """
    模型类，封装了Transformer模型并添加了保存和加载功能。
    """
    def __init__(
            self,
            config,
            logger,
            input_dim,
            output_dim,
    ):
        """
        初始化模型类。
        :param config: Config, 配置类，包含模型配置参数。
        :param logger: Logger, 日志记录器。
        :param input_dim: int, 输入数据的维度。
        :param output_dim: int, 输出数据的维度。
        """
        super(Model, self).__init__(input_dim, output_dim, num_heads=config.NUM_HEADS,
                                    hidden_dim=config.HIDDEN_DIM, num_layers=config.NUM_LAYERS,
                                    dropout=config.DROPOUT_RATE)
        self.config = config
        self.logger = logger

    def save(
            self,
            path,
    ):
        """
        保存模型到指定路径。
        :param path: str, 保存模型的文件路径。
        """
        torch.save(self.state_dict(), path)
        self.logger.info(f"模型存储路径：{path}")

    def load(
            self,
            path,
    ):
        """
        从指定路径加载模型。
        :param path: str, 加载模型的文件路径。
        """
        self.load_state_dict(torch.load(path, map_location=device))
        self.logger.info(f"从 {path} 加载模型")
        return self


def train_model(
        model,
        dataloader,
        logger,
):
    """
    训练模型。
    :param model: Model, 要训练的模型。
    :param dataloader: DataLoader, 用于加载训练数据的DataLoader。
    :param logger: Logger, 日志记录器。
    :return: Model, 训练好的模型。
    """
    model.to(device)
    criterion = nn.MSELoss().to(device)
    optimizer = optim.Adam(model.parameters(), lr=TaskConfig.LR)
    epochs = TaskConfig.EPOCHS
    for epoch in range(epochs):
        progress_bar = tqdm(dataloader, desc=f'Epoch {epoch + 1}/{epochs}')
        for batch, (inputs, targets) in enumerate(progress_bar):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            progress_bar.set_postfix({'Loss': loss.item()})
            logger.info(f"Epoch {epoch + 1}, Batch {batch + 1}, Loss: {loss.item()}")
        torch.cuda.empty_cache()
    return model


def predict(
        model,
        test_data,
        scaler,
):
    """
    对测试数据进行推理。
    :param model: Model, 推理模型。
    :param test_data: torch.Tensor, 测试数据。
    :param scaler: MinMaxScaler, 缩放器，对推理结果进行逆变换。
    :return: pd.DataFrame, 包含时间戳和流量预测值的DataFrame。
    """
    test_data = test_data.to(device)
    with torch.no_grad():
        model.eval()
        predict_tensor = model(test_data)
    predict_array = predict_tensor.cpu().numpy().reshape(-1, 1)
    predict_array_rescaled = scaler.inverse_transform(predict_array).reshape(-1, )
    timestamps = generate_timestamps('2019/3/1 00:00:00', '2019/6/18 23:00:00', 60)
    return pd.DataFrame({'时间': timestamps, '流量': predict_array_rescaled})
